///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/tei.cpp
///
/// \brief   implementation file for two electron integral functions
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <algorithm>

#include <formic/ham/tei.h>
#include <formic/lapack/interface.h>
#include <formic/numeric/numeric.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   expands compactly stored two electron integrals into dense storage
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_cmp   the two electron integrals in compact storage
/// \param[out]   tei_dns   on exit, the two electron integrals in dense storage
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_compact_to_dense(const int n, const S * const tei_cmp, S * const tei_dns) {

  for (int c = 0; c < n ; c++)
  for (int d = 0; d <= c; d++)
  for (int a = 0; a < n ; a++)
  for (int b = 0; b <= a; b++)
  {
    const S integ = tei_cmp[ formic::tei_compound_index(a,b,c,d) ];
    tei_dns[a+n*(b+n*(c+n*d))] = integ;
    tei_dns[a+n*(b+n*(d+n*c))] = integ;
    tei_dns[b+n*(a+n*(c+n*d))] = integ;
    tei_dns[b+n*(a+n*(d+n*c))] = integ;
  }

}

template void formic::tei_compact_to_dense(const int n, const double * const tei_cmp, double * const tei_dns);
template void formic::tei_compact_to_dense(const int n, const std::complex<double> * const tei_cmp, std::complex<double> * const tei_dns);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   converts two electron integrals from dense storage to compact storage
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_dns   the two electron integrals in dense storage
/// \param[out]   tei_cmp   on exit, the two electron integrals in compact storage
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_dense_to_compact(const int n, const S * const tei_dns, S * const tei_cmp) {

  std::fill(tei_cmp, tei_cmp + formic::tei_compact_array_length(n), formic::zero(S()));
  for (int d = 0; d < n; d++)
  for (int c = 0; c < n; c++)
  for (int b = 0; b < n; b++)
  for (int a = 0; a < n; a++) {
    double x = 1.0;
    x *= ( a == b ? 1.0 : 0.5 );
    x *= ( c == d ? 1.0 : 0.5 );
    x *= ( std::min(a,b) == std::min(c,d) && std::max(a,b) == std::max(c,d) ? 1.0 : 0.5 );
    tei_cmp[ formic::tei_compound_index(a,b,c,d) ] += x * tei_dns[a+n*(b+n*(c+n*d))];
  }

}

template void formic::tei_dense_to_compact(const int n, const double * const tei_dns, double * const tei_cmp);
template void formic::tei_dense_to_compact(const int n, const std::complex<double> * const tei_dns, std::complex<double> * const tei_cmp);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   rotates the two electron integral indices from  p,q,r,s  to  q,r,s,p
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_in    the dense two electron integrals in p,q,r,s order
/// \param[out]   tei_out   the dense two electron integrals in q,r,s,p order
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_rotate_indices_left(const int n, const S * const tei_in, S * const tei_out) {
  const int n2 = n*n;
  const int n3 = n*n*n;
  for (int s = 0; s < n; s++)
  for (int p = 0; p < n; p++)
    // copy all q,r for these values of p and s
    formic::xcopy(n2, tei_in + n3*s + p, n, tei_out + n3*p + n2*s, 1);
}

template void formic::tei_rotate_indices_left(const int n, const double * const tei_in, double * const tei_out);
template void formic::tei_rotate_indices_left(const int n, const std::complex<double> * const tei_in, std::complex<double> * const tei_out);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   rotates the two electron integral indices from  p,q,r,s  to  s,p,q,r
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_in    the dense two electron integrals in p,q,r,s order
/// \param[out]   tei_out   the dense two electron integrals in s,p,q,r order
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_rotate_indices_right(const int n, const S * const tei_in, S * const tei_out) {
  const int n2 = n*n;
  const int n3 = n*n*n;
  //for (int s = 0; s < n; s++)
  //for (int r = 0; r < n; r++)
  //  // copy all p,q for these values of r and s
  //  formic::xcopy(n2, tei_in + n3*s + n2*r, 1, tei_out + n3*r + s, n);
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
  for (int r = 0; r < n; r++)
  for (int s = 0; s < n; s++)
    tei_out[ s + n*p + n2*q + n3*r ] = tei_in[ p + n*q + n2*r + n3*s ];
    //tei_out[ q + n*r + n2*s + n3*p ] = tei_in[ p + n*q + n2*r + n3*s ];
}

template void formic::tei_rotate_indices_right(const int n, const double * const tei_in, double * const tei_out);
template void formic::tei_rotate_indices_right(const int n, const std::complex<double> * const tei_in, std::complex<double> * const tei_out);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reorders the two electron integral indices from  p,q,r,s  to  q,p,r,s
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_in    the dense two electron integrals in p,q,r,s order
/// \param[out]   tei_out   the dense two electron integrals in q,p,r,s order
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_swap_first_two_indices(const int n, const S * const tei_in, S * const tei_out) {

  const int two_n = 2*n;
  const int n2 = n*n;
  const int n3 = n*n*n;

  //for (int s = 0; s < n; s++)
  //for (int r = 0; r < n; r++) {
  //  const int offset = n3 * s + n2 * r;
  //  const S * x = tei_in  + offset;
  //        S * y = tei_out + offset;
  //  for (int q = 0; q < n; q++)
  //  for (int p = 0; p < n; p++)
  //    y[q+p*n] = x[p+q*n];
  //}

  for (int s = 0; s < n; s++)
  for (int r = 0; r < n; r++) {
    const int offset = n3 * s + n2 * r;
    const S * const x = tei_in  + offset;
          S * const y = tei_out + offset;
    for (int q = 0; q < n; q++) {
            S * yb = y + q;
      const S * xb = x + q*n;
      const S * xe = xb + n;
      if ( n % 2 ) {
        yb[0] = xb[0];
        yb += n;
        xb += 1;
      }
      while (xb != xe) {
        yb[0] = xb[0];
        yb[n] = xb[1];
        yb += two_n;
        xb += 2;
      }
    }
  }

}

template void formic::tei_swap_first_two_indices(const int n, const double * const tei_in, double * const tei_out);
template void formic::tei_swap_first_two_indices(const int n, const std::complex<double> * const tei_in, std::complex<double> * const tei_out);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reorders the two electron integral indices from  p,q,r,s  to  p,q,s,r
/// 
/// \param[in]    n         number of orbitals
/// \param[in]    tei_in    the dense two electron integrals in p,q,r,s order
/// \param[out]   tei_out   the dense two electron integrals in p,q,s,r order
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::tei_swap_last_two_indices(const int n, const S * const tei_in, S * const tei_out) {

  const int n2 = n*n;
  const int n3 = n*n*n;

  for (int s = 0; s < n; s++)
  for (int r = 0; r < n; r++) {
    const S * const x = tei_in  + n3 * s + n2 * r;
          S * const y = tei_out + n3 * r + n2 * s;
    formic::xcopy(n2, x, 1, y, 1);
  }

}

template void formic::tei_swap_last_two_indices(const int n, const double * const tei_in, double * const tei_out);
template void formic::tei_swap_last_two_indices(const int n, const std::complex<double> * const tei_in, std::complex<double> * const tei_out);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Extracts the occ-unocc-occ-unocc, aaaa slice of the two electron integrals, in which
///          elements with repeated occupied and/or unoccupied indices are set to zero.
/// 
/// \param[in]    no          number of occupied orbitals
/// \param[in]    nv          number of unoccupied orbitals
/// \param[in]    ao2s        occupied alpha sites, size no
/// \param[in]    bo2s        occupied  beta sites, size no
/// \param[in]    au2s        unoccupied alpha sites, size nv
/// \param[in]    bu2s        unoccupied  beta sites, size nv
/// \param[in]    tei_dense   the dense two electron integrals, size (no+nv)^4
/// \param[out]   tei_slice   on exit, the ovov aaaa slice, size no*no*nv*nv
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::tei_get_ovov_aaaa_slice(const int no,
                                                       const int nv,
                                                       const int * const ao2s,
                                                       const int * const bo2s,
                                                       const int * const au2s,
                                                       const int * const bu2s,
                                                       const S * const tei_dense,
                                                             S * const tei_slice) {

      const int nt = no + nv;
      std::fill(tei_slice, tei_slice + no*no*nv*nv, formic::zero(S()));
      for (int b = 0; b < nv; b++)
      for (int j = 0; j < no; j++)
      for (int a = 0; a < b  ; a++)
      for (int i = 0; i < j  ; i++) {
        const S integ1 = tei_dense[ ao2s[i] + nt * ( au2s[a] + nt * ( ao2s[j] + nt * au2s[b] ) ) ];
        const S integ2 = tei_dense[ ao2s[j] + nt * ( au2s[a] + nt * ( ao2s[i] + nt * au2s[b] ) ) ];
        tei_slice[ i + no * ( a + nv * ( j + no * b ) ) ] = integ1;
        tei_slice[ j + no * ( b + nv * ( i + no * a ) ) ] = integ1;
        tei_slice[ j + no * ( a + nv * ( i + no * b ) ) ] = integ2;
        tei_slice[ i + no * ( b + nv * ( j + no * a ) ) ] = integ2;
      }

}

template void formic::tei_get_ovov_aaaa_slice(const int no,
                                              const int nv,
                                              const int * const ao2s,
                                              const int * const bo2s,
                                              const int * const au2s,
                                              const int * const bu2s,
                                              const double * const tei_dense,
                                                    double * const tei_slice);
template void formic::tei_get_ovov_aaaa_slice(const int no,
                                              const int nv,
                                              const int * const ao2s,
                                              const int * const bo2s,
                                              const int * const au2s,
                                              const int * const bu2s,
                                              const std::complex<double> * const tei_dense,
                                                    std::complex<double> * const tei_slice);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Extracts the occ-unocc-occ-unocc, aabb slice of the two electron integrals.
/// 
/// \param[in]    no          number of occupied orbitals
/// \param[in]    nv          number of unoccupied orbitals
/// \param[in]    ao2s        occupied alpha sites, size no
/// \param[in]    bo2s        occupied  beta sites, size no
/// \param[in]    au2s        unoccupied alpha sites, size nv
/// \param[in]    bu2s        unoccupied  beta sites, size nv
/// \param[in]    tei_dense   the dense two electron integrals, size (no+nv)^4
/// \param[out]   tei_slice   on exit, the ovov aabb slice, size no*no*nv*nv
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::tei_get_ovov_aabb_slice(const int no,
                                                       const int nv,
                                                       const int * const ao2s,
                                                       const int * const bo2s,
                                                       const int * const au2s,
                                                       const int * const bu2s,
                                                       const S * const tei_dense,
                                                             S * const tei_slice) {

      const int nt = no + nv;
      for (int b = 0; b < nv; b++)
      for (int j = 0; j < no; j++)
      for (int a = 0; a < nv; a++)
      for (int i = 0; i < no; i++) {
        const S integral = tei_dense[ ao2s[i] + nt * ( au2s[a] + nt * ( bo2s[j] + nt * bu2s[b] ) ) ];
        tei_slice[ i + no * ( a + nv * ( j + no * b ) ) ] = integral;
      }

}

template void formic::tei_get_ovov_aabb_slice(const int no,
                                              const int nv,
                                              const int * const ao2s,
                                              const int * const bo2s,
                                              const int * const au2s,
                                              const int * const bu2s,
                                              const double * const tei_dense,
                                                    double * const tei_slice);
template void formic::tei_get_ovov_aabb_slice(const int no,
                                              const int nv,
                                              const int * const ao2s,
                                              const int * const bo2s,
                                              const int * const au2s,
                                              const int * const bu2s,
                                              const std::complex<double> * const tei_dense,
                                                    std::complex<double> * const tei_slice);
